import torch
import torchvision
from utils.train import Accumulator
import utils.object_detection as object_detection
import torch.utils.data
import matplotlib.pyplot as plot
import torch.nn.functional as f


class Train:
    def __init__(self, **kwargs):
        self.sequential = kwargs['sequential']
        self.optimizer = kwargs['optimizer']
        self.scheduler = kwargs['scheduler']
        self.train_iter = kwargs['train_iter']
        self.epochs = kwargs['epochs']
        self.device = kwargs['device']
        self.accumulator = kwargs['accumulator']

        self.c = kwargs['c']
        self.b = kwargs['b']

        self.s = kwargs['s']
        self.m = kwargs['m']
        self.sizes = kwargs['sizes']
        self.ratios = kwargs['ratios']

    def train(self):
        self.sequential.train(mode=True)
        self.sequential.to(self.device)
        accuracies = []
        losses = []
        m = self.m
        s = self.s

        anchor_boxes = object_detection. \
            multi_box(m, m, sizes=self.sizes, ratios=self.ratios, device=self.device). \
            repeat(self.train_iter.batch_size, 1, 1)  # size: (batch_size, m*m*5(b+c), 4)

        cell_offsets_w = torch.arange(0, m) * (1 / m)
        cell_offsets_w = cell_offsets_w.repeat(m, 1)
        cell_offsets_h = torch.arange(0, m) * (1 / m)
        cell_offsets_h = cell_offsets_h.reshape(m, 1).repeat(1, m)
        cell_offsets = torch.stack([cell_offsets_w, cell_offsets_h], dim=2). \
            repeat((1, 1, s)). \
            reshape(m * m * s, 2). \
            repeat(self.train_iter.batch_size, 1, 1). \
            reshape(self.train_iter.batch_size, m * m * s, 2). \
            to(self.device)
        anchor_boxes[:, :, :] = object_detection.bbox_position_cast(anchor_boxes.reshape(-1, 4), True).reshape(
            self.train_iter.batch_size, -1, 4)
        anchor_boxes.requires_grad = False
        cell_offsets.requires_grad = False

        for e in range(self.epochs):
            for x, y in self.train_iter:
                self.optimizer.zero_grad()
                x, y = x.to(self.device), y.to(self.device)
                y = y.unsqueeze(1)
                truths, masks = yolo_target_anchor_boxes(y, m, s, self.device)

                y_pred = self.sequential(x)
                y_pred = y_pred.permute(0, 2, 3, 1)
                y_pred = y_pred.reshape(y_pred.shape[0], -1, self.b + self.c)

                l = loss(y_pred, truths, masks, cell_offsets, anchor_boxes, y[:, :, 1:5], m, s).mean()
                l.backward()
                self.optimizer.step()

                ev, ev1 = self.train_accuracy(y_pred, truths, masks,
                                              anchor_boxes, offsets=cell_offsets)
                print("\r", f"loss: {l.item()}", f"evaluation: {ev.item()} , {ev1.item()}",
                      end='')
                self.accumulator.collect(l.item(), ev)
            # self.scheduler.step()
            accuracies.append(self.accumulator.ap())
            avg_loss = self.accumulator.al()
            print('     epoch ' + str(e) + ' avg loss: ' + str(avg_loss) + ' accuracy: ' + str(self.accumulator.ap()))
            self.accumulator.reset()
            losses.append(avg_loss)

    def train_accuracy(self, predicts, truths, masks, anchor_boxes, offsets):
        cn = predicts.shape[0]
        predicts_confidences = predicts[:, :, 0:1]  # 置信度
        predicts_boxes = predicts[:, :, 1:5]
        predicts_class_scores = predicts[:, :, 5:6]
        truths_class_scores = truths[:, :, 0:1]
        truths_boxes = truths[:, :, 1:5]
        target_confidences = torch.sigmoid(predicts_confidences[masks])

        target_truths_boxes = truths_boxes[masks].unsqueeze(1)  # 真实框

        predicts_boxes_clone = decode(predicts_boxes=predicts_boxes.detach(), offsets=offsets.detach(),
                                      anchor_boxes=anchor_boxes.detach(), m=self.m)
        target_predict_boxes = predicts_boxes_clone[:, :, 0:4][masks].unsqueeze(1)  # 预测框

        iou_table = object_detection.iou(target_truths_boxes.reshape(-1, 4), target_predict_boxes.reshape(-1, 4),
                                         mode='union')
        indices = f.one_hot(torch.argmax(iou_table.reshape(cn, -1), dim=1), num_classes=self.s).bool().view(-1)
        return target_confidences[indices].mean(), iou_table[indices].mean()
        # return object_detection.mean_ap(target_confidences, target_truths_boxes, target_predict_boxes, cn)


def decode(predicts_boxes, offsets, anchor_boxes, m):
    cn = predicts_boxes.shape[0]
    offsets_cn = offsets[0:cn, :, :]
    anchor_boxes_cn = anchor_boxes[0:cn, :, :]
    predicts_boxes_clone = torch.clone(predicts_boxes)

    predicts_boxes_clone[:, :, 0:2] = torch.sigmoid(predicts_boxes_clone[:, :, 0:2]) * 1. / m + offsets_cn
    predicts_boxes_clone[:, :, 2:4] = anchor_boxes_cn[:, :, 2:4] * torch.exp(predicts_boxes_clone[:, :, 2:4])
    predicts_boxes_clone[:, :, :] = object_detection.bbox_position_cast(predicts_boxes_clone.reshape(-1, 4),
                                                                        False).reshape(cn, -1,
                                                                                       4)

    return predicts_boxes_clone


def loss(predicts, truths, masks, offsets, anchor_boxes, labels_boxes, m, s):
    cn = predicts.shape[0]  # 记录当前batch_size
    predicts_confidences = predicts[:, :, 0:1]  # 置信度
    predicts_boxes = predicts[:, :, 1:5]
    predicts_class_scores = predicts[:, :, 5:6]
    truths_class_scores = truths[:, :, 0:1]
    truths_boxes = truths[:, :, 1:5]

    anchor_boxes_cn = anchor_boxes[0:cn, :, :]
    anchor_boxes_cn_w_h = anchor_boxes_cn[:, :, 2:4]
    offsets_cn = offsets[0:cn, :, :]

    truths_boxes_xy_wh = object_detection.bbox_position_cast(truths_boxes.reshape(-1, 4), True).reshape(cn, -1, 4)

    predicts_boxes_1 = torch.clone(predicts_boxes)

    predicts_boxes_1[:, :, 0:2] = torch.sigmoid(predicts_boxes_1[:,:,0:2]) * 1. / m + offsets_cn
    predicts_boxes_1[:, :, 2:4] = anchor_boxes_cn_w_h * torch.exp(predicts_boxes_1[:, :, 2:4])
    predicts_boxes_1[:, :, :] = object_detection.bbox_position_cast(predicts_boxes_1.reshape(-1, 4),
                                                                    False).reshape(cn, -1,
                                                                                   4)



    mse_loss = torch.nn.MSELoss(reduction='none')

    iou_tables = object_detection.iou(predicts_boxes_1.view(-1, 4), truths_boxes.view(-1, 4), mode='union').view(cn, -1)

    no_obj_masks = torch.bitwise_not(masks)
    l1 = 1 * no_obj_masks * mse_loss(
        torch.zeros(predicts_confidences.shape, device=predicts_confidences.device),
        torch.sigmoid(predicts_confidences)).mean(dim=2) + 5 * masks * mse_loss(iou_tables.unsqueeze(2),
                                                                        torch.sigmoid(predicts_confidences)).sum(dim=2)
    l2 = 1 * masks * mse_loss(truths_class_scores, torch.sigmoid(predicts_class_scores)).sum(dim=2)
    l3 = 1 * masks * mse_loss(torch.sigmoid(predicts_boxes[:, :, 0:2]) * 1. / m + offsets_cn,
                                      truths_boxes_xy_wh[:, :, 0:2]).sum(
        dim=2)

    l4 = 1 * masks * mse_loss(predicts_boxes[:, :, 2:4],
                                      torch.log(truths_boxes_xy_wh[:, :, 2:4] / anchor_boxes_cn_w_h + 1e-16)).sum(
        dim=2)
    # with torch.no_grad():
    #     print(l1.sum(dim=1).mean(), l2.sum(dim=1).mean(), l3.sum(dim=1).mean(), l4.sum(dim=1).mean())

    return l1.sum(dim=1) + l2.sum(dim=1) + l3.sum(dim=1) + l4.sum(dim=1)


def yolo_target_anchor_boxes(labels, m, s, device):
    boxes, masks = [], []

    n = m * m * s

    for i in range(labels.shape[0]):
        label = labels[i, :, :]
        label_boxes = label[:, 1:5]
        label_boxes_x_y_w_h = object_detection.bbox_position_cast(label_boxes, True)
        cell_j = torch.floor(label_boxes_x_y_w_h[:, 0] * m).long()
        cell_i = torch.floor(label_boxes_x_y_w_h[:, 1] * m).long()
        cell_index = cell_i * m + cell_j
        truths = torch.zeros((n, 5), dtype=torch.float32, device=device).reshape(m * m, s, -1)
        truths[cell_index, :, :] = label.unsqueeze(0).repeat(1, s, 1)
        truths = truths.reshape(m * m * s, -1)
        mask = torch.zeros(n, dtype=torch.bool, device=device).reshape(m * m, s)
        mask[cell_index, :] = True
        mask = mask.reshape(m * m * s)
        boxes.append(truths)
        masks.append(mask)

    boxes = torch.stack(boxes)
    masks = torch.stack(masks)

    # 真实框，每个格子包含的目标类，掩码
    return boxes, masks
